// internal/console/graphics_proxy.go
package console

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/sirupsen/logrus"
	"libvirt.org/go/libvirtxml"

	"virt-manager-go/internal/domain"
)

// GraphicsProxy VNC/SPICE 图形代理
type GraphicsProxy struct {
	connMgr  domain.ConnectionManager
	logger   *logrus.Logger
	upgrader websocket.Upgrader

	// 代理会话管理
	sessions   map[string]*graphicsSession
	sessionsMu sync.RWMutex

	// 令牌管理（用于认证）
	tokens   map[string]*proxyToken
	tokensMu sync.RWMutex
}

type graphicsSession struct {
	id          string
	vmName      string
	sessionType string // vnc, spice
	wsConn      *websocket.Conn
	tcpConn     net.Conn
	ctx         context.Context
	cancel      context.CancelFunc
	logger      *logrus.Logger
}

type proxyToken struct {
	token     string
	connID    string
	vmName    string
	proxyType string
	expiresAt time.Time
}

func NewGraphicsProxy(connMgr domain.ConnectionManager, logger *logrus.Logger) *GraphicsProxy {
	return &GraphicsProxy{
		connMgr: connMgr,
		logger:  logger,
		upgrader: websocket.Upgrader{
			ReadBufferSize:  16384,
			WriteBufferSize: 16384,
			CheckOrigin: func(r *http.Request) bool {
				return true // 生产环境需要验证 Origin
			},
		},
		sessions: make(map[string]*graphicsSession),
		tokens:   make(map[string]*proxyToken),
	}
}

// GetVNCInfo 获取 VNC 连接信息
func (p *GraphicsProxy) GetVNCInfo(ctx context.Context, connID, vmName string) (*domain.ProxyInfo, error) {
	conn, err := p.connMgr.GetConnection(connID)
	if err != nil {
		return nil, err
	}

	dom, err := conn.LookupDomainByName(vmName)
	if err != nil {
		return nil, fmt.Errorf("VM not found: %w", err)
	}
	defer dom.Free()

	// 解析 XML 获取 VNC 配置
	xmlData, err := dom.GetXMLDesc(0)
	if err != nil {
		return nil, fmt.Errorf("failed to get VM XML: %w", err)
	}

	domainConfig := &libvirtxml.Domain{}
	if err := domainConfig.Unmarshal(xmlData); err != nil {
		return nil, fmt.Errorf("failed to parse VM XML: %w", err)
	}

	// 查找 VNC 图形设备
	var vncPort int
	var vncListen string
	for _, gfx := range domainConfig.Devices.Graphics {
		if gfx.VNC != nil {
			vncPort = gfx.VNC.Port
			if len(gfx.VNC.Listeners) > 0 && gfx.VNC.Listeners[0].Address != nil {
				vncListen = gfx.VNC.Listeners[0].Address.Address
			}
			break
		}
	}

	if vncPort == 0 {
		return nil, fmt.Errorf("VNC not configured for this VM")
	}

	// 生成访问令牌
	token := p.generateToken(connID, vmName, "vnc")

	proxyInfo := &domain.ProxyInfo{
		ProxyURL:  fmt.Sprintf("%s:%d", vncListen, vncPort),
		WSPath:    fmt.Sprintf("/api/v1/console/vnc/%s/%s?token=%s", connID, vmName, token),
		Token:     token,
		ExpiresAt: time.Now().Add(1 * time.Hour),
	}

	return proxyInfo, nil
}

// GetSPICEInfo 获取 SPICE 连接信息
func (p *GraphicsProxy) GetSPICEInfo(ctx context.Context, connID, vmName string) (*domain.ProxyInfo, error) {
	conn, err := p.connMgr.GetConnection(connID)
	if err != nil {
		return nil, err
	}

	dom, err := conn.LookupDomainByName(vmName)
	if err != nil {
		return nil, fmt.Errorf("VM not found: %w", err)
	}
	defer dom.Free()

	xmlData, err := dom.GetXMLDesc(0)
	if err != nil {
		return nil, fmt.Errorf("failed to get VM XML: %w", err)
	}

	domainConfig := &libvirtxml.Domain{}
	if err := domainConfig.Unmarshal(xmlData); err != nil {
		return nil, fmt.Errorf("failed to parse VM XML: %w", err)
	}

	// 查找 SPICE 图形设备
	var spicePort int
	var spiceListen string
	for _, gfx := range domainConfig.Devices.Graphics {
		if gfx.Spice != nil {
			spicePort = gfx.Spice.Port
			if len(gfx.Spice.Listeners) > 0 && gfx.Spice.Listeners[0].Address != nil {
				spiceListen = gfx.Spice.Listeners[0].Address.Address
			}
			break
		}
	}

	if spicePort == 0 {
		return nil, fmt.Errorf("SPICE not configured for this VM")
	}

	token := p.generateToken(connID, vmName, "spice")

	proxyInfo := &domain.ProxyInfo{
		ProxyURL:  fmt.Sprintf("spice://%s:%d", spiceListen, spicePort),
		WSPath:    fmt.Sprintf("/api/v1/console/spice/%s/%s?token=%s", connID, vmName, token),
		Token:     token,
		ExpiresAt: time.Now().Add(1 * time.Hour),
	}

	return proxyInfo, nil
}

// HandleVNCWebSocket 处理 VNC WebSocket 连接
func (p *GraphicsProxy) HandleVNCWebSocket(w http.ResponseWriter, r *http.Request, connID, vmName string) error {
	// 验证令牌
	token := r.URL.Query().Get("token")
	if !p.validateToken(token, connID, vmName, "vnc") {
		http.Error(w, "Invalid or expired token", http.StatusUnauthorized)
		return fmt.Errorf("invalid token")
	}

	// 升级到 WebSocket
	ws, err := p.upgrader.Upgrade(w, r, nil)
	if err != nil {
		return fmt.Errorf("failed to upgrade to websocket: %w", err)
	}

	p.logger.WithFields(logrus.Fields{
		"connection": connID,
		"vm":         vmName,
		"remote":     r.RemoteAddr,
		"type":       "vnc",
	}).Info("VNC WebSocket connection established")

	// 获取 VNC 连接信息
	proxyInfo, err := p.GetVNCInfo(r.Context(), connID, vmName)
	if err != nil {
		ws.Close()
		return err
	}

	// 连接到 VNC 服务器
	vncConn, err := net.DialTimeout("tcp", proxyInfo.ProxyURL[6:], 10*time.Second)
	if err != nil {
		ws.Close()
		return fmt.Errorf("failed to connect to VNC server: %w", err)
	}

	// 创建代理会话
	ctx, cancel := context.WithCancel(r.Context())
	session := &graphicsSession{
		id:          fmt.Sprintf("vnc-%s-%d", vmName, time.Now().Unix()),
		vmName:      vmName,
		sessionType: "vnc",
		wsConn:      ws,
		tcpConn:     vncConn,
		ctx:         ctx,
		cancel:      cancel,
		logger:      p.logger,
	}

	// 注册会话
	p.sessionsMu.Lock()
	p.sessions[session.id] = session
	p.sessionsMu.Unlock()

	// 清理会话
	defer func() {
		p.sessionsMu.Lock()
		delete(p.sessions, session.id)
		p.sessionsMu.Unlock()

		vncConn.Close()
		ws.Close()
		cancel()

		p.logger.WithField("session_id", session.id).Info("VNC session closed")
	}()

	// 启动双向代理
	return p.proxyBidirectional(session)
}

// HandleSPICEWebSocket 处理 SPICE WebSocket 连接
func (p *GraphicsProxy) HandleSPICEWebSocket(w http.ResponseWriter, r *http.Request, connID, vmName string) error {
	// 验证令牌
	token := r.URL.Query().Get("token")
	if !p.validateToken(token, connID, vmName, "spice") {
		http.Error(w, "Invalid or expired token", http.StatusUnauthorized)
		return fmt.Errorf("invalid token")
	}

	// 升级到 WebSocket
	ws, err := p.upgrader.Upgrade(w, r, nil)
	if err != nil {
		return fmt.Errorf("failed to upgrade to websocket: %w", err)
	}

	p.logger.WithFields(logrus.Fields{
		"connection": connID,
		"vm":         vmName,
		"remote":     r.RemoteAddr,
		"type":       "spice",
	}).Info("SPICE WebSocket connection established")

	// 获取 SPICE 连接信息
	proxyInfo, err := p.GetSPICEInfo(r.Context(), connID, vmName)
	if err != nil {
		ws.Close()
		return err
	}

	// 连接到 SPICE 服务器
	spiceConn, err := net.DialTimeout("tcp", proxyInfo.ProxyURL[8:], 10*time.Second)
	if err != nil {
		ws.Close()
		return fmt.Errorf("failed to connect to SPICE server: %w", err)
	}

	// 创建代理会话
	ctx, cancel := context.WithCancel(r.Context())
	session := &graphicsSession{
		id:          fmt.Sprintf("spice-%s-%d", vmName, time.Now().Unix()),
		vmName:      vmName,
		sessionType: "spice",
		wsConn:      ws,
		tcpConn:     spiceConn,
		ctx:         ctx,
		cancel:      cancel,
		logger:      p.logger,
	}

	// 注册会话
	p.sessionsMu.Lock()
	p.sessions[session.id] = session
	p.sessionsMu.Unlock()

	// 清理会话
	defer func() {
		p.sessionsMu.Lock()
		delete(p.sessions, session.id)
		p.sessionsMu.Unlock()

		spiceConn.Close()
		ws.Close()
		cancel()

		p.logger.WithField("session_id", session.id).Info("SPICE session closed")
	}()

	// 启动双向代理
	return p.proxyBidirectional(session)
}

// proxyBidirectional 双向代理数据传输
func (p *GraphicsProxy) proxyBidirectional(session *graphicsSession) error {
	errCh := make(chan error, 2)

	// WebSocket -> TCP
	go p.wsToTCP(session, errCh)

	// TCP -> WebSocket
	go p.tcpToWS(session, errCh)

	// 等待任一方向出错或上下文取消
	select {
	case err := <-errCh:
		if err != nil && err != io.EOF {
			session.logger.WithError(err).Warn("Graphics proxy error")
		}
		return err
	case <-session.ctx.Done():
		return session.ctx.Err()
	}
}

// wsToTCP WebSocket 到 TCP 的数据转发
func (p *GraphicsProxy) wsToTCP(session *graphicsSession, errCh chan<- error) {
	defer session.cancel()

	for {
		select {
		case <-session.ctx.Done():
			return
		default:
		}

		msgType, data, err := session.wsConn.ReadMessage()
		if err != nil {
			if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
				errCh <- fmt.Errorf("websocket read error: %w", err)
			}
			return
		}

		if msgType != websocket.BinaryMessage {
			continue
		}

		if _, err := session.tcpConn.Write(data); err != nil {
			errCh <- fmt.Errorf("tcp write error: %w", err)
			return
		}
	}
}

// tcpToWS TCP 到 WebSocket 的数据转发
func (p *GraphicsProxy) tcpToWS(session *graphicsSession, errCh chan<- error) {
	defer session.cancel()

	buffer := make([]byte, 32768) // 32KB 缓冲区

	for {
		select {
		case <-session.ctx.Done():
			return
		default:
		}

		n, err := session.tcpConn.Read(buffer)
		if err != nil {
			if err != io.EOF {
				errCh <- fmt.Errorf("tcp read error: %w", err)
			}
			return
		}

		if n == 0 {
			continue
		}

		if err := session.wsConn.WriteMessage(websocket.BinaryMessage, buffer[:n]); err != nil {
			errCh <- fmt.Errorf("websocket write error: %w", err)
			return
		}
	}
}

// generateToken 生成访问令牌
func (p *GraphicsProxy) generateToken(connID, vmName, proxyType string) string {
	tokenBytes := make([]byte, 32)
	rand.Read(tokenBytes)
	token := hex.EncodeToString(tokenBytes)

	p.tokensMu.Lock()
	p.tokens[token] = &proxyToken{
		token:     token,
		connID:    connID,
		vmName:    vmName,
		proxyType: proxyType,
		expiresAt: time.Now().Add(1 * time.Hour),
	}
	p.tokensMu.Unlock()

	// 启动令牌清理
	go p.cleanupExpiredToken(token)

	return token
}

// validateToken 验证令牌
func (p *GraphicsProxy) validateToken(token, connID, vmName, proxyType string) bool {
	p.tokensMu.RLock()
	defer p.tokensMu.RUnlock()

	t, exists := p.tokens[token]
	if !exists {
		return false
	}

	if time.Now().After(t.expiresAt) {
		return false
	}

	return t.connID == connID && t.vmName == vmName && t.proxyType == proxyType
}

// cleanupExpiredToken 清理过期令牌
func (p *GraphicsProxy) cleanupExpiredToken(token string) {
	time.Sleep(1 * time.Hour)

	p.tokensMu.Lock()
	delete(p.tokens, token)
	p.tokensMu.Unlock()
}

// GetActiveSessions 获取活跃会话
func (p *GraphicsProxy) GetActiveSessions() []*graphicsSession {
	p.sessionsMu.RLock()
	defer p.sessionsMu.RUnlock()

	sessions := make([]*graphicsSession, 0, len(p.sessions))
	for _, s := range p.sessions {
		sessions = append(sessions, s)
	}
	return sessions
}
